#include "addcmul_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 64;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
  AddcmulTilingData tiling;
  uint32_t sizeofdatatype;
  uint32_t totalLengthAligned;
  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  auto socVersion = ascendcPlatform.GetSocVersion();
  uint64_t ub_size;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  auto aivNum = ascendcPlatform.GetCoreNumAiv();  // vector core num  1
  
  uint32_t totalLength = context->GetInputTensor(0)->GetShapeSize();
  uint32_t x1Length = context->GetInputTensor(1)->GetShapeSize();
  if(x1Length != totalLength){
    tiling.set_x1Length(x1Length);
  }else{
    tiling.set_x1Length(0);
  }
  uint32_t x2Length = context->GetInputTensor(2)->GetShapeSize();
  if(x2Length != totalLength){
    tiling.set_x2Length(x2Length);
  }else{
    tiling.set_x2Length(0);
  }
  auto dt = context->GetInputTensor(0)->GetDataType();
  uint32_t typeKey = 0;
  uint32_t data_num = 5; // kernel中用到了多少组buffer和buf
  if (dt == 1) { // half
    sizeofdatatype = 2;
    tiling.set_typeKey(1);
  }else if (dt == 0) { // float
    sizeofdatatype = 4;
    tiling.set_typeKey(0);
  }else if (dt == 2) { // int8
    sizeofdatatype = 1;
    data_num = 12;
    tiling.set_typeKey(2);
  }else if (dt == 3) { // int32
    sizeofdatatype = 4;
    tiling.set_typeKey(3);
  }

  uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
  //核内拆分，策略是尽可能的填满ub_size，最后一包单独处理，
  uint32_t ub_block_num =
      ((ub_size) / BLOCK_SIZE / data_num) * 7 / 20;  // ub_block在Ascend C中不能全部被用来作为输入输出，给了13/20系数。
  uint32_t tile_num;

  if (ub_block_num % 2 != 0) {
    ub_block_num = ub_block_num - 1;
  }

  // 1.输入向量满足32字节对齐
  if (totalLength % ALIGN_NUM != 0) {  //不对齐，先32位对齐
    totalLengthAligned =
        ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
  } else {
    totalLengthAligned = totalLength;
  }

  context->SetBlockDim(1);

  uint32_t tileLength = 0;
  uint32_t lasttileLength = 0;

  tile_num = totalLengthAligned / ALIGN_NUM / ub_block_num; // 简单获取tile_num的大小
  if ((totalLengthAligned / ALIGN_NUM) % ub_block_num == 0 ||
      tile_num == 0) {  //对齐之后，可以均分到各个tile，或不足一个tile
    if (tile_num == 0) {
      tile_num = 1;
    } 
    if (totalLengthAligned < ub_block_num * ALIGN_NUM) {
      tileLength = totalLengthAligned;
      lasttileLength = tileLength;
    } else {
      tileLength = ub_block_num * ALIGN_NUM;
      lasttileLength = tileLength;
    }
  } else {  //满足字节对齐，核内不能均分
    tile_num = tile_num + 1;
    tileLength = ub_block_num * ALIGN_NUM;
    lasttileLength = totalLengthAligned - (tile_num - 1) * tileLength;
  }
  printf("%u\t\t%u******%u*******\n", tile_num, tileLength, lasttileLength);
  tiling.set_blockLength(totalLengthAligned);
  tiling.set_tileNum(tile_num);
  tiling.set_tileLength(tileLength);
  tiling.set_lasttileLength(lasttileLength);

  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                      context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  size_t* currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = 0;
  return ge::GRAPH_SUCCESS;

}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Addcmul : public OpDef {
public:
    explicit Addcmul(const char* name) : OpDef(name)
    {
        this->Input("input_data")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("value")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Addcmul);
}
